library(R2jags)
library(xlsx)
library(openxlsx)
library(spatstat)
library(dplyr)
library(readr)
require(R2wd)
library(LaplacesDemon) # for Geweke.Diagnostic()
library(ggplot2)

options(max.print=999999999)

#memory.limit(10000);memory.size(max=T)

getparams<-function(jagssum,parameters=c('mu','sigma','diff','disc','beta'),questions=questions){
  for (n in parameters){v<-data.frame(jagssum[grepl(n,dimnames(jagssum)[[1]]),"mean"])
    s<-data.frame(jagssum[grepl(n,dimnames(jagssum)[[1]]),"sd"])
    for (w in c(parameters,'\\[','\\]')){row.names(v)<-gsub(w,"",row.names(v));row.names(s)<-gsub(w,"",row.names(s))}
    v <- cbind(rownames(v),v);s<-cbind(rownames(s),s)
    if (n=='beta'){v<-data.frame(v[,1:2]);s<-data.frame(s[,1:2]);colnames(v)<-c('n','mean.beta');colnames(s)<-c('n','sd.beta');pq<-left_join(v,s) 
    } else {if (n=="mu"|n=="sigma"){colnames(v)<-c('t',paste0("mean.",n));colnames(s)<-c('t',paste0("sd.",n));pq<-left_join(v,s)
      } else {if (n=="diff"){colnames(v)<-c('qn',paste0('mean.',n));colnames(s)<-c('qn',paste0('sd.',n))
          if (nrow(v)==info$nquest){v$qn<-as.character(v$qn);s$qn<-as.character(s$qn)
          } else {v$b<-substr(v$qn,(nchar(as.character(v$qn))+1)-1,nchar(as.character(v$qn)))
            s$b<-substr(s$qn,(nchar(as.character(s$qn))+1)-1,nchar(as.character(s$qn)))
            row.names(v)<-c(1:nrow(v));row.names(s)<-c(1:nrow(s))
            v$qn<-gsub(",1|,2|,3","",v$qn)
            s$qn<-gsub(",1|,2|,3","",s$qn)}
        } else {colnames(v)<-c('qn',paste0('mean.',n));colnames(s)<-c('qn',paste0('sd.',n))
          v$qn<-as.character(v$qn);s$qn<-as.character(s$qn)}
        pq<-left_join(v,s);pq<-left_join(pq,questions);pq$question<-as.character(pq$question)}}
    return(pq)}}

setwd(dirname(rstudioapi::getActiveDocumentContext()$path)) # set wd to source file location

'%!in%' <- Negate('%in%')

## JAGS CODE ######################################################################
ms_model="data{for (i in 1:len){opinionr[i,2]<-round(support[i]*samplesize[i]/100)
		opinionr[i,1]<-round(neutral[i]*samplesize[i]/100)		}}
model{for (i in 1:len){for (j in 1:2){
		opinionr[i,j]~dbin(p[i,j],samplesize[i])      
		p[i,j]~dbeta(alpha[i,j], beta)              
		alpha[i,j]<--beta*m[i,j]/(m[i,j]-1)         
		m[i,j]<-phi(x[i,j])                   
	  	x[i,j]<-(mu[t[i]]-diff[qn[i],j])/sqrt((disc[qn[i]])^2+(sigma[t[i]])^2)}}
sd<-0.1
tau.beta<-1/(sd*sd)
for (i in 1:nquest){tau.disc[i]<-1/(sd*sd)
for (j in 1:2){tau.diff[i,j]<-1/(sd*sd)}}
beta~dnorm(mean.beta,tau.beta) T(0,) # Posterior distribution as prior
mu[begin]~dnorm(0,1)
sigma_mu~dunif(0,10)
for (i in begin:end){sigma[i]~dunif(0,10)} 
for (i in begin:(end-1)){mu_raw[i]~dnorm(0,1)}
for (i in (begin+1):end){mu[i]~dnorm(mu[i-1]+(sigma_mu*mu_raw[i-1]),0.001)} 
for (i in 1:nquest){disc[i]~dnorm(mean.disc[i],tau.disc[i]) T(0,) 
	for (j in 1:2){diff0[i,j]~dnorm(mean.diff[i+nquest*(j-1)],tau.diff[i,j])}
	diff[i,1:2]<-sort(diff0[i,])}}"
## Save to temp file
fileConn=file("ms_model.temp")
writeLines(ms_model,fileConn)
close(fileConn)
###################################################################################


outdir<-'irt_estimates/';dir.create(file.path(getwd(),outdir),showWarnings=F)
inpdir<-'raw_indicators/'

mainsheet<-read.table(paste0(inpdir,'EU_irt.txt'),sep=';',header=T)  # Read in data
mainsheet<-subset(mainsheet,as.numeric(gsub(' S1| S2','',mainsheet$t))>=1973)
mainsheet$t<-as.numeric(mainsheet$t)
mainsheet$n<-NULL
mainsheet<-suppressMessages(left_join(mainsheet,mainsheet %>% group_by(question) %>% tally() %>% as.data.frame()))
names(mainsheet)[names(mainsheet)=='n']<-'asked'  
mainsheet<-subset(mainsheet,mainsheet$asked>1)

policy<-list('jags','asylum','market','econom','social') # Delete or keep accordingly

for (p in policy){sheet<-mainsheet;sheet$policy<-sheet$thetopic;if (p!='jags'){sheet<-sheet[sheet$policy==p,]}
  rownames(sheet) <- 1:nrow(sheet);sheet$question<-factor(sheet$question) # Reset row numbers and update levels of q names after subsetting data
  sheet$qn<-as.numeric(sheet$question);sheet<-sheet[order(sheet$qn,sheet$t),] # Sort by q number and time
  questions<-unique(sheet[c("qn","question")]);questions$question<-as.character(questions$question);questions$qn<-as.character(questions$qn)
  # Loop over countries
  country<-list('BE','DE','FR','IT','LU','NL','DK','IE','GB','GR','ES','PT','AT','FI','SE','CY','CZ','EE','HU','LV','LT','MT','PL','SI','SK','BG','RO','HR')
      for (c in country){print(paste0('Country: ',c,'. Policy: ',p))
    wb<-paste0(outdir,c,'_irt_estimates.xlsx');filename<-paste0(inpdir,c,'_irt.txt') # Store workbook and filename
    ms<-read.table(filename,sep=';',header=T)  # Read in data
    ms<-subset(ms,as.numeric(gsub(' S1| S2','',ms$t))>=1973)
    ms$t<-as.numeric(ms$t);ms$asked<-NULL
    ms<-suppressMessages(left_join(ms,ms %>% group_by(question) %>% tally() %>% as.data.frame()))
    names(ms)[names(ms)=='n']<-'asked'  
    ms<-subset(ms,ms$asked>1)
    ms$policy<-ms$thetopic
    # Filter out observations according to policy domain
    if (p!='jags'){ms<-ms[ms$policy==p,]}

    # Read in priors for question parameters stored in txt
    a.constant<-read.table(paste0(outdir,'model_parameters/discrimination_EU_',p,'.txt'),sep="\t",header=T)
    b.constant<-read.table(paste0(outdir,'model_parameters/difficulty_EU_',p,'.txt'),sep="\t",header=T)
    
    # Remove questions lacking EU-level parameters
    ms<-ms[ms$question %in% c(a.constant$question),]
    
    beta.constant<-read.table(paste0(outdir,'model_parameters/beta_EU_',p,'.txt'),sep="\t",header=T)
    if (!"b" %in% colnames(b.constant)){b.constant$b<-"1"} # Add column if only 1 b parameter estimated
    a.prior<-unique(left_join(ms,a.constant,by=c('question'))[c('mean.disc','sd.disc','question')])
    b.prior<-unique(left_join(ms,b.constant,by=c('question'))[c('mean.diff','sd.diff','question','b')])
    b.prior$b<-as.numeric(b.prior$b);colnames(a.prior)<-c('mean.disc','sd.disc','question');colnames(b.prior)<-c('mean.diff','sd.diff','question','b')
    b.prior<-unique(b.prior[order(b.prior$b,b.prior$question),][c('question','mean.diff','sd.diff','b')]);b.prior$question<-NULL
    a.prior<-unique(a.prior[order(a.prior$question),][c('question','mean.disc','sd.disc')]);a.prior$question<-NULL
    beta.prior<-beta.constant[c('mean.beta','sd.beta')]
    
    # Prepare data
    rownames(ms)<-1:nrow(ms);ms$question<-factor(ms$question) # Reset row numbers and update "levels" of q names if subsetting data
    ms$qn<-as.numeric(ms$question);ms<-ms[order(ms$qn,ms$t),] # Adjust var type and sort
    questions.ms<-unique(ms[c("qn","question")]);questions.ms$question<-as.character(questions.ms$question);questions.ms$qn<-as.character(questions.ms$qn);ms$question<-as.character(ms$question)
    ms<-left_join(ms,questions,by=c('question'));ms$qn.y<-NULL;names(ms)[names(ms)=='qn.x']<-'qn' #  Merge ms data with q data and rename columns
    
    info<-list(len=nrow(ms),nquest=length(unique(ms$question)),begin=min(ms$t),end=max(ms$t));datams<-c(ms,info,a.prior,b.prior,beta.prior)
    start_time<-Sys.time();print(paste0('Using ',info$nquest,' questions and ',info$len,' observations. ',start_time))
    invisible(gc());jagsfitms<-jags.parallel(data=datams,parameters.to.save=c('mu','sigma',"diff","disc","beta"),
                                             n.burnin=10000,n.iter=100000,n.thin=2,n.chains=3,
                                             model.file="ms_model.temp",jags.seed=123)
    duration<-Sys.time()-start_time;print(duration)
    fit.mcmc.ms<-as.mcmc(jagsfitms)
    print(gelman.diag(fit.mcmc.ms,autoburnin=F)$mpsrf)

    resultsms<-jagsfitms$BUGSoutput$summary
    
    sink(file=paste0(outdir,'model_summaries/gelman_diagnostic_',c,'_',p,'.txt'));print(gelman.diag(fit.mcmc.ms,autoburnin=F));print('LIST OF QUESTIONS');print(questions.ms);sink()
    sink(file=paste0(outdir,'model_summaries/geweke_diagnostic_',c,'_',p,'.txt'));print(Geweke.Diagnostic(fit.mcmc.ms));print('LIST OF QUESTIONS');print(questions.ms);sink()
    sink(file=paste0(outdir,'model_summaries/summary_model_',c,'_',p,'.txt'));print(duration);print(jagsfitms);sink()
    
    for (par in c('mu','sigma','diff','disc','beta')){p.func<-getparams(jagsfitms$BUGSoutput$summary,par,questions.ms);assign(par,p.func)};remove(jagsfitms);print(" ")

    write.table(beta,file=paste0(outdir,'model_parameters/beta_',c,'_',p,'.txt'),row.names=F,quote=F,sep="\t")
    write.table(diff,file=paste0(outdir,'model_parameters/difficulty_',c,'_',p,'.txt'),row.names=F,quote=F,sep="\t")
    write.table(disc,file=paste0(outdir,'model_parameters/discrimination_',c,'_',p,'.txt'),row.names=F,quote=F,sep="\t")
    
    df<-merge(mu,sigma);df$t<-as.numeric(as.character(df$t));df<-df[order(df$t),]
    if ((nrow(b.prior))==info$nquest){diff<-b.prior} else {j=as.character(max(b.prior$b));diff<-subset(b.prior,b==j)}
    for (r in 1:nrow(df)){
      df$Opinion[r]<-pnorm((df$mean.mu[r]-mean(diff$mean.diff))/(mean(disc$mean.disc)^2+df$mean.sigma[r]^2)^.5,0,1)*100
      df$Opinion_minus2sd[r]<-pnorm(((df$mean.mu[r]-2*df$sd.mu[r])-mean(diff$mean.diff))/(mean(disc$mean.disc)^2+df$mean.sigma[r]^2)^.5,0,1)*100
      df$Opinion_plus2sd[r]<-pnorm(((df$mean.mu[r]+2*df$sd.mu[r])-mean(diff$mean.diff))/(mean(disc$mean.disc)^2+df$mean.sigma[r]^2)^.5,0,1)*100
      df$Opinion_minussigma[r]<-df$Opinion[r]-(mean(disc$mean.disc)^2+df$mean.sigma[r]^2)^.5
      df$Opinion_plussigma[r]<-df$Opinion[r]+(mean(disc$mean.disc)^2+df$mean.sigma[r]^2)^.5}
    chart<-ggplot(df, aes(x = t, y = Opinion, group = 1)) + geom_line(col='black') + xlab('Year') + ggtitle(paste0(p,' ',c)) +
      geom_ribbon(aes(ymin = Opinion_minus2sd, ymax = Opinion_plus2sd), alpha = 0.1) + theme_bw()
    #pdf(paste0(outdir,'irt_plots/',c,'_',p,'.pdf'),width=6.5,height=4);chart;dev.off()
    
    if (!file.exists(wb)){workbook<-createWorkbook();addWorksheet(workbook,p);writeData(workbook,sheet=p,x=df,rowNames=F);saveWorkbook(workbook,wb,overwrite=T)} else {workbook<-loadWorkbook(wb);if (p %in% names(workbook)){removeWorksheet(workbook,p);addWorksheet(workbook,p)} else {addWorksheet(workbook,p)};writeData(workbook,sheet=p,x=df,rowNames=F)
      saveWorkbook(workbook,wb,overwrite=T)}}}
